#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
"""
@File    :   SSH.py
@Time    :   2024-11-27 09:59
@Author  :   坐公交也用券
@Version :   1.0
@Contact :   faith01238@hotmail.com
@Homepage : https://liumou.site
@Desc    :   当前文件作用
"""
from argparse import ArgumentParser
from sys import exit
import nmap
from loguru import logger


class SSH:
	def __init__(self, ip, mask=24, port=22):
		"""
		初始化网络扫描器

		:param ip: 设备的IP地址
		:param mask: 子网掩码，默认为24
		:param port: 要扫描的端口，默认为22
		"""
		self.ip = ip
		self.mask = mask

		# 创建 self.portscanner 对象
		self.nm = nmap.PortScanner()

		# 定义要扫描的网络段
		self.network = f"{ip}/{mask}"

		# 定义要扫描的端口，SSH 默认端口是 22
		self.ports = port

	def start(self):
		"""
		将开放的SSH扫描结果写入文件。

		此方法遍历扫描结果中的所有主机，检查指定端口是否开放。
		如果开放，则写入主机信息和SSH版本信息到文件；如果没有开放，则写入相应信息。
		"""
		# 扫描网络
		self.nm.scan(hosts=str(self.network), ports=str(self.ports), arguments='-sV')  # -sV 用于服务版本检测
		close = open(f"results_close_{self.ip}-{self.mask}.txt", "w")
		open_file = open(f"results_open_{self.ip}-{self.mask}.txt", "w")
		try:
			# 保存扫描结果到文件
			for host in self.nm.all_hosts():
				# 获取主机IP地址
				host_ip = self.nm[host].hostname()
				logger.info(f"Host: {host_ip} ({host})")
				ssh_info = self.nm[host]['tcp'][self.ports]
				logger.info(ssh_info)
				# 检查指定端口是否开放
				if str(ssh_info['state']).lower() == 'open'.lower():
					# 获取并写入SSH版本信息
					ssh_version = ssh_info['version']
					logger.info(f"{host} - SSH_Version: {ssh_version}")
					open_file.write(f"{host} {ssh_version}\n")
				else:
					# 如果端口未开放，写入相应信息
					logger.warning(f"SSH not open: {host}")
					close.write(f"{host}\n")
		except Exception as e:
			logger.error(e)
			exit(3)


if __name__ == '__main__':
	# 传参参数及说明
	arg = ArgumentParser(description='OpenSSH扫描结果', prog="SSH")
	arg.add_argument("-i", "--ip", help="IP地址,必须,例如: 192.168.1.1", type=str, required=True)
	arg.add_argument("-m", "--mask", help="子网掩码,默认: 24", type=int, default=24)
	arg.add_argument("-p", "--port", help="端口,默认: 22", type=int, default=22)
	# 定义SSH版本要求
	arg.add_argument("-v", "--version", help="指定SSH版本,默认: 9.9", type=str, default="9.9")
	args = arg.parse_args()
	logger.info(args)
	logger.info("开始扫描")
	ssh = SSH(args.ip, args.mask, args.port)
	ssh.start()
